from omegaconf import DictConfig

import numpy as np
from torch.utils.data import Dataset


class BaseImageDataset(Dataset):
    """
    Base dataset that loads images and can be used for other datasets
    """
    def __init__(
            self,
            cfg: DictConfig,
            mode: str
            ):
        if mode == 'train':
            self.data_imaging = np.load(cfg.data_imaging_train, mmap_mode='c')
            self.data_imaging_es_timesteps = np.load(cfg.data_es_timesteps_imaging_train, mmap_mode='c')
        elif mode == 'val':
            self.data_imaging = np.load(cfg.data_imaging_val, mmap_mode='c')
            self.data_imaging_es_timesteps = np.load(cfg.data_es_timesteps_imaging_val, mmap_mode='c')
        elif mode == 'test':
            self.data_imaging = np.load(cfg.data_imaging_test, mmap_mode='c')
            self.data_imaging_es_timesteps = np.load(cfg.data_es_timesteps_imaging_test, mmap_mode='c')

    def __len__(self):
        return len(self.data_imaging)

    def __getitem__(self, index: int) -> np.ndarray:
        return self.data_imaging[index], self.data_imaging_es_timesteps[index]
    